#include "ToolsCapability.h"
#include "CapabilityKeyword.h"
#include "CapabilityMethod.h"

#include "message_serialization/JsonRpcErrorCode.h"
#include "protocol/JsonRpcMessageSerializer.h"

#include <QLatin1String>
#include <QJsonDocument>
#include <QJsonParseError>
#include <QDebug>
#include <memory>

static const QLatin1String TOOLS_LIST("tools/list");
static const QLatin1String TOOLS_CALL("tools/call");

ToolsCapability::ToolsCapability()
    : mSerializer(std::make_unique<JsonRpcMessageSerializer>())
{
    mCapability[Capability::LIST_CHANGED] = false;
}

QString ToolsCapability::prefix() const
{
    return "tools";
}

QJsonObject ToolsCapability::capability() const
{
    return mCapability;
}

void ToolsCapability::setCapability(const QString& capability, const bool enable)
{
    if (capability == Capability::LIST_CHANGED) {
        mCapability[Capability::LIST_CHANGED] = enable;
    }
}

QJsonArray ToolsCapability::list() const
{
    return mToolsList;
}

bool ToolsCapability::addMethod(const CapabilityMethod& method)
{
    const QJsonObject& methodDesc = method.describe;
    
    // Validate required fields based on ToolsMethodBuilder format
    if (!methodDesc.contains("name") || !methodDesc["name"].isString()) {
        qWarning() << "Tool method missing or invalid name field";
        return false;
    }
    
    QString methodName = methodDesc["name"].toString();
    if (methodName.isEmpty()) {
        qWarning() << "Tool method name cannot be empty";
        return false;
    }
    
    // Check if method already exists
    if (mMethodMap.contains(methodName)) {
        qWarning() << "Tool method already exists:" << methodName;
        return false;
    }
    
    // Validate optional description field
    if (methodDesc.contains("description") && !methodDesc["description"].isString()) {
        qWarning() << "Tool method description must be a string";
        return false;
    }
    
    // Validate optional title field
    if (methodDesc.contains("title") && !methodDesc["title"].isString()) {
        qWarning() << "Tool method title must be a string";
        return false;
    }
    
    // Validate inputSchema if present
    if (methodDesc.contains("inputSchema")) {
        if (!methodDesc["inputSchema"].isObject()) {
            qWarning() << "Tool method inputSchema must be a JSON object";
            return false;
        }
        
        QJsonObject inputSchema = methodDesc["inputSchema"].toObject();
        if (!inputSchema.contains("type") || inputSchema["type"].toString() != "object") {
            qWarning() << "Tool method inputSchema must have type 'object'";
            return false;
        }
    }
    
    // Validate outputSchema if present
    if (methodDesc.contains("outputSchema")) {
        if (!methodDesc["outputSchema"].isObject()) {
            qWarning() << "Tool method outputSchema must be a JSON object";
            return false;
        }
        
        QJsonObject outputSchema = methodDesc["outputSchema"].toObject();
        if (!outputSchema.contains("type") || outputSchema["type"].toString() != "object") {
            qWarning() << "Tool method outputSchema must have type 'object'";
            return false;
        }
    }
    
    if (method.task == nullptr) {
        qWarning() << "Tool method missing callback function";
        return false;
    }

    // Store in JSON array for listing
    mToolsList.append(methodDesc);
    
    // Extract parameter information for validation and store in mMethodMap
    ToolsCapability::Method methodDef;
    methodDef.name = methodName;
    
    // Extract parameter types and required list from inputSchema
    if (methodDesc.contains("inputSchema")) {
        QJsonObject inputSchema = methodDesc["inputSchema"].toObject();
        
        if (inputSchema.contains("properties") && inputSchema["properties"].isObject()) {
            QJsonObject properties = inputSchema["properties"].toObject();
            
            for (auto it = properties.begin(); it != properties.end(); ++it) {
                QString paramName = it.key();
                QJsonValue paramValue = it.value();
                
                if (paramValue.isObject()) {
                    QJsonObject paramObj = paramValue.toObject();
                    if (paramObj.contains("type") && paramObj["type"].isString()) {
                        QString paramType = paramObj["type"].toString();
                        methodDef.paramTypeMap[paramName] = paramType;
                    }
                }
            }
        }
        
        // Extract required parameters
        if (inputSchema.contains("required") && inputSchema["required"].isArray()) {
            QJsonArray requiredArray = inputSchema["required"].toArray();
            for (const QJsonValue& requiredValue : qAsConst(requiredArray)) {
                if (requiredValue.isString()) {
                    methodDef.requiredList.append(requiredValue.toString());
                }
            }
        }
    }
    
    // Store method definition for validation
    mMethodMap[methodName] = methodDef;
    
    // Store the task for execution
    mTaskMap[methodName] = method.task;
    
    return true;
}

void ToolsCapability::clearMethod()
{
    mToolsList = QJsonArray();
    mMethodMap.clear();
    mTaskMap.clear();
}

QJsonObject ToolsCapability::doRequest(const QJsonObject& request)
{
    if (!mSerializer) {
        qWarning() << "No serializer set for ToolsCapability";
        return QJsonObject();
    }

    QVariant id = request.value("id").toVariant();
    QString methodName = request["method"].toString();

    if (methodName == TOOLS_LIST) {
        QJsonObject tools{{"tools", this->list()}};
        return mSerializer->serializeSuccessResponse(id, tools);
    }
    else if (methodName == TOOLS_CALL) {
        QString msg;
        int errorCode = this->checkMethodValidate(request, &msg);
        if (errorCode != 0) {
            qWarning() << msg;
            return mSerializer->serializeErrorResponse(id, errorCode, QJsonValue(msg));
        }

        // Extract method name and arguments from request
        QJsonObject params = request["params"].toObject();
        QString toolName = params["name"].toString();
        QJsonObject arguments = params["arguments"].toObject();

        // Execute the task with the provided arguments
        try {
            QJsonObject executionResult = mTaskMap[toolName]->execute(arguments);
            return mSerializer->serializeSuccessResponse(id, executionResult);
        }
        catch (const std::exception& e) {
            QString msg = QString("Tool execution failed: %1").arg(e.what());
            qWarning() << msg;
            return mSerializer->serializeErrorResponse
                    (id, RpcError::SERVER_ERROR, QJsonValue(msg));
        }
    }
    else {
        QString msg = QString("Method not found: %1").arg(methodName);
        qWarning() << msg;
        return mSerializer->serializeErrorResponse(
                    id, RpcError::METHOD_NOT_FOUND, QJsonValue(QString(msg)));
    }
}

int ToolsCapability::checkMethodValidate(const QJsonObject& request, QString* msg) const
{
    QJsonObject params = request["params"].toObject();
    QString methodName = params["name"].toString();
    if (!mMethodMap.contains(methodName)) {
        if (msg) *msg = QString("Unknown tool: %1").arg(methodName);
        return RpcError::INVALID_PARAMS;
    }

    auto methodDef = mMethodMap.value(methodName);
    QJsonObject arguments = params["arguments"].toObject();
    QStringList keys = arguments.keys();
    for (const auto& required : qAsConst(methodDef.requiredList)) {
        if (!keys.contains(required)) {
            if (msg) *msg = QString("Missing required argument: %1").arg(required);
            return RpcError::INVALID_PARAMS;
        }
    }

    const auto& paramMap = methodDef.paramTypeMap;
    for (const auto& arg : qAsConst(keys)) {
        if (!paramMap.contains(arg)) {
            continue;
        }
        QString type = paramMap[arg];
        QJsonValue data = arguments[arg];
        if      (type == "string" && data.isString()) continue;
        else if (type == "number" && data.isDouble()) continue;
        else if (type == "object" && data.isObject()) continue;
        else if (type == "array"  && data.isArray())  continue;

        if (msg) *msg = QString("Not a matching argument, expect: %1").arg(type);
        return RpcError::INVALID_PARAMS;
    }

    return 0;
}

void ToolsCapability::setSerializer(std::unique_ptr<IMessageSerializer> serializer)
 {
    if (serializer) {
        mSerializer = std::move(serializer);
    }
 }
